package flow

import (
	"sync"
	"time"
)

var TokenBucketRedis *TokenBucket = NewTokenBucket(10, 10)

/* Token Bucket algorithm for flow control*/
type TokenBucket struct {
	rate         int64
	cap          int64
	lock         sync.Mutex
	tokens       int64 //Current number of tokens
	lastPutToken int64 //last time of putting token, second timestamps
}

func NewTokenBucket(rate int64, cap int64) *TokenBucket {
	return &TokenBucket{
		rate:         rate,
		cap:          cap,
		tokens:       0,
		lastPutToken: time.Now().Unix(),
	}
}

/* Set an existing TokenBucket */
func (tb *TokenBucket) Set(rate int64, cap int64) {
	tb.lock.Lock()
	defer tb.lock.Unlock()

	tb.rate = rate
	tb.cap = cap
	tb.tokens = 0
	tb.lastPutToken = time.Now().Unix()
}

/* Allow or not */
func (tb *TokenBucket) Allow() bool {
	tb.lock.Lock()
	defer tb.lock.Unlock()

	tb.addTokenRealtime()
	if tb.tokens > 0 {
		tb.tokens = tb.tokens - 1
		return true
	} else {
		return false
	}
}

/* Add tokens to bucket by real time, called by TokenBucket.Allow() */
func (tb *TokenBucket) addTokenRealtime() {
	now := time.Now().Unix()
	tokens := tb.tokens + (now-tb.lastPutToken)*tb.rate

	if tokens > tb.cap {
		tb.tokens = tb.cap
	} else {
		tb.tokens = tokens
	}

	tb.lastPutToken = now
}

func (tb *TokenBucket) AllowWithTimes(n int) bool {
	var allow bool
	for i := 0; i < n; i++ {
		if tb.Allow() {
			allow = true
			break
		}
		time.Sleep(200 * time.Millisecond)
	}

	return allow
}
